#include "nvidia/cuda_common.h"

#include <iostream>

namespace autodrive {
namespace perception {

CUDATimer::CUDATimer(cudaStream_t stream) noexcept(false) : stream_{stream} {
  auto cuRet{cudaEventCreate(&start_)};
  check_cuda_error(cuRet);

  cuRet = cudaEventCreate(&stop_);
  check_cuda_error(cuRet);

  cuRet = cudaEventRecord(start_, stream_);
  check_cuda_error(cuRet);
}

float CUDATimer::Elapsed() const noexcept(false) {
  float time{0};

  auto cuRet{cudaEventRecord(stop_, stream_)};
  check_cuda_error(cuRet);

  cuRet = cudaEventSynchronize(stop_);
  check_cuda_error(cuRet);

  cuRet = cudaEventElapsedTime(&time, start_, stop_);
  check_cuda_error(cuRet);

  return time;
}

CUDATimer::~CUDATimer() {
  try {
    auto cuRetStart{cudaEventDestroy(start_)};

    auto cuRetStop{cudaEventDestroy(stop_)};

    if (cudaSuccess != cuRetStart) {
      std::cerr << "[CUDA] cudaEventDestroy start error, error code="
                << cuRetStart << '\n';
    }

    if (cudaSuccess != cuRetStop) {
      std::cerr << "[CUDA] cudaEventDestroy stop error, error code="
                << cuRetStop << '\n';
    }

  } catch (const std::exception& e) {
    std::cerr << "[CUDA] cudaEventDestroy error, " << e.what() << '\n';
  } catch (...) {
    std::cerr << "[CUDA] cudaEventDestroy error, Unexpected" << '\n';
  }
}

void set_cuda_device(std::int32_t device, std::ostream& info) noexcept(false) {
  cudaError_t cudaRet{cudaSuccess};
  cudaRet = cudaSetDevice(device);
  check_cuda_error(cudaRet);

  cudaDeviceProp properties;
  cudaRet = cudaGetDeviceProperties(&properties, device);
  check_cuda_error(cudaRet);

  info << "=== Device Information ===" << std::endl;
  info << "Selected Device: " << properties.name << std::endl;
  info << "Compute Capability: " << properties.major << "." << properties.minor
       << std::endl;
  info << "SMs: " << properties.multiProcessorCount << std::endl;
  info << "Compute Clock Rate: " << properties.clockRate / 1000000.0F << " GHz"
       << std::endl;
  info << "Device Global Memory: " << (properties.totalGlobalMem >> 20)
       << " MiB" << std::endl;
  info << "Shared Memory per SM: "
       << (properties.sharedMemPerMultiprocessor >> 10) << " KiB" << std::endl;
  info << "Memory Bus Width: " << properties.memoryBusWidth << " bits"
       << " (ECC " << (properties.ECCEnabled != 0 ? "enabled" : "disabled")
       << ")" << std::endl;
  info << "Memory Clock Rate: " << properties.memoryClockRate / 1000000.0F
       << " GHz" << std::endl;
  info << "Can Map Host Memory: " << properties.canMapHostMemory << std::endl;
  info << "=== Device Information End===" << std::endl;
}

void pin_cuda_memory(void* host_addr, std::int32_t size,
                     void** dev_addr) noexcept(false) {
  if (nullptr == host_addr) {
    throw std::runtime_error("host address is nullptr");
  }

  auto cuRet{cudaHostRegister(host_addr, size, cudaHostRegisterMapped)};
  check_cuda_error(cuRet);

  cuRet = cudaHostGetDevicePointer(dev_addr, host_addr, 0);
  check_cuda_error(cuRet);
}

CUDAEventPtr make_cuda_event() noexcept(false) {
  auto deleter{[](cudaEvent_t* e) {
    if (nullptr != e) {
      auto cuRet{cudaEventDestroy(*e)};
      check_cuda_error(cuRet);
    }
    delete e;
  }};

  auto cudaEventPtr{CUDAEventPtr{new cudaEvent_t, deleter}};
  auto cuRet{cudaEventCreate(cudaEventPtr.get())};
  if (cudaSuccess != cuRet) {
    cudaEventPtr.reset();
  }
  check_cuda_error(cuRet);
  return cudaEventPtr;
}

CUDAPinnedMemoryPtr make_cuda_pinned_mem(std::int32_t size) noexcept(false) {
  void* hostAddr;
  auto cuRet{cudaMallocHost(&hostAddr, size)};
  check_cuda_error(cuRet);

  void* devAddr;
  cuRet = cudaHostGetDevicePointer(&devAddr, hostAddr, 0);
  check_cuda_error(cuRet);

  auto deleter{[](CUDAPinnedMemory* mem) {
    if (mem != nullptr) {
      if (mem->host != nullptr) {
        auto cuRet{cudaFreeHost(mem->host)};
        check_cuda_error(cuRet);
      }

      mem->host = nullptr;
      mem->dev = nullptr;
    }
  }};

  return CUDAPinnedMemoryPtr{
      new CUDAPinnedMemory(CUDAPinnedMemory{.host = hostAddr, .dev = devAddr}),
      deleter};
}

}  // namespace perception
}  // namespace autodrive